model = My_VGG16(class_num=2)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.Adam(model.parameters(), 0.001)
for epoch in range(50):
    for step,(x, y) in enumerate(train_loader): 
        logits = model(x)
        y_true = torch.tensor(data = y,dtype=torch.long,device=device)
        loss = criterion(input=logits, target=y_true)
        # 准确率
        pred = torch.argmax(logits, axis=1)
        correct=torch.tensor(torch.eq(pred, y_true), dtype=torch.float32)
        acc = torch.mean(correct)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if step % 10 == 0:
            print("epoch:", epoch, "step:", step,
                  'loss:', float(loss), "acc:",acc.numpy())
torch.save(model.state_dict(), 'catdog_vgg.pt') 
